##########################################################################################

setwd("~/20220915_gastric_multiple/dna_combinePublic/scripts/evolutionTime")

##########################################################################################

library(rjags)
library(tidyverse)
library(ggmcmc)
library(kableExtra)
library(magrittr)
library(readxl)

##########################################################################################

input_file <- "~/20220915_gastric_multiple/dna_combinePublic/scripts/evolutionTime/ipmn-timing-master/data/Timing_Metrics.xlsx"
out_path <- "~/20220915_gastric_multiple/dna_combinePublic/scripts/evolutionTime"
code_path <- "~/20220915_gastric_multiple/dna_combinePublic/scripts/evolutionTime/ipmn-timing-master"

##########################################################################################

source(paste0(code_path, "/code/functions.R"))
jag_file <- paste0( code_path , "/code/clock.jag"  )

##########################################################################################

unixFriendly <- function(x){
  x <- colnames(x)
  x <- tolower(x)
  x <- gsub(" ", "_", x)
  x
}

##########################################################################################

dat <- read_excel(input_file,
                  sheet=2) %>%
  "["(-nrow(.), ) %>%
  set_colnames(unixFriendly(.)) %>%
  mutate(additional=mutations_ca-mutations_hg_ipmn)
avgN <- dat %>%
  summarize(mean=mean(additional)) %>%
  "[["(1)


##########################################################################################

mu1 <- round(avgN/5, 2)
mu2 <- round(avgN/4, 2)
mu3 <- round(avgN/2, 2)

##########################################################################################
## Model 1: `r mu1` mutations/year
y <- dat$additional
m1 <- jags.model(jag_file,
                data=list(y=y, N = length(y), mu=mu1),
                inits=list(T=rep(5, length(y))),
                n.adapt = 2000, n.chains = 3)
jags_samples <- coda.samples(m1,
                             variable.names="T",
                             n.iter=2000)
jags_df <- ggs(jags_samples)
##ggs_traceplot(jags_df)
summaries.mu1 <- jags_df %>%
  "["(grep("T", .$Parameter), ) %>%
  group_by(Parameter) %>%
  summarize(lower=quantile(value, 0.05),
            median=quantile(value, 0.5),
            mean=mean(value),
            upper=quantile(value, 0.95)) %>%
  mutate(mu=mu1)


## Model 2: `r mu2` mutations/year
m2 <- jags.model(jag_file,
                data=list(y=y, N = length(y), mu=mu2),
                inits=list(T=rep(5, length(y))),
                n.adapt = 2000, n.chains = 3)
jags_samples <- coda.samples(m2,
                             variable.names="T",
                             n.iter=2000)
jags_df <- ggs(jags_samples)
summaries.mu2 <- jags_df %>%
  "["(grep("T", .$Parameter), ) %>%
  group_by(Parameter) %>%
  summarize(lower=quantile(value, 0.05),
            median=quantile(value, 0.5),
            mean=mean(value),
            upper=quantile(value, 0.95)) %>%
  mutate(mu=mu2)

## Model 3: `r mu3` mutations/year
m3 <- jags.model(jag_file,
                data=list(y=y, N = length(y), mu=mu3),
                inits=list(T=rep(5, length(y))),
                n.adapt = 2000, n.chains = 3)
jags_samples <- coda.samples(m3,
                             variable.names="T",
                             n.iter=2000)
jags_df <- ggs(jags_samples)
summaries.mu3 <- jags_df %>%
  "["(grep("T", .$Parameter), ) %>%
  group_by(Parameter) %>%
  summarize(lower=quantile(value, 0.05),
            median=quantile(value, 0.5),
            mean=mean(value),
            upper=quantile(value, 0.95)) %>%
  mutate(mu=mu3)

##########################################################################################
## Posterior summaries
summaries <- list(summaries.mu1,
                  summaries.mu2,
                  summaries.mu3) %>%
  do.call(rbind, .) %>%
  mutate(id=rep(dat$case, 3),
         id=factor(id, levels=dat$case))
mus <- rep(paste(c(mu1, mu2, mu3), "mutations/year"), each=nrow(dat)) %>%
  factor(., levels=unique(.))

##########################################################################################
## Timing estimates faceted by assumed mutation rate
facetlabels <- setNames(levels(mus), c(mu1, mu2, mu3))
p <- ggplot(summaries, aes(ymin=lower, ymax=upper, x=id, xend=id)) +
  geom_point(aes(y=median), shape=21, color="gray") +
  geom_errorbar(width=0.2) +
  coord_flip() +
  facet_wrap(~mu, labeller=labeller(mu=facetlabels)) +
  theme(panel.background=element_rect(fill="white", color="black")) +
  xlab("mutations/year") + ylab("Years")

out_name <- paste0( out_path , "/timing_estimates.mutationrate.pdf" )
ggsave( out_name , p )

##########################################################################################
## Timing estimates faceted by sample id
p <- ggplot(summaries, aes(ymin=lower, ymax=upper, x=mu, xend=mu)) +
  geom_point(aes(y=median), shape=21, color="gray") +
  geom_errorbar(width=0.2) +
  coord_flip() +
  facet_wrap(~id) +
  theme(panel.background=element_rect(fill="white", color="black")) +
  xlab("mutations/year") + ylab("Years")

out_name <- paste0( out_path , "/timing_estimates.sampleid.pdf" )
ggsave( out_name , p )

##########################################################################################
## Tabled values
result <- summaries %>%
  select(id, lower, median, mean, upper, mu) %>%
  mutate(lower=round(lower, 2),
         median=round(median, 2),
         mean=round(mean, 2),
         upper=round(upper, 2)) %>%
  set_colnames(c("Id", "Lower (years)", "Median (years)",
"Mean (years)", "Upper (years)", "Mutations/year")) %>%
  kable("html", escape=FALSE) %>%
  kable_styling(bootstrap_options=c("striped", "hover", "condensed"),
                full_width=FALSE)